#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 内存Cookie提取工具

import frida
import sys
import re
import json
import os
import argparse
from datetime import datetime

class CookieExtractor:
    def __init__(self, output_dir="cookies_extracted"):
        self.cookies_by_user = {}
        self.total_cookies = 0
        self.output_dir = output_dir
        # 创建保存结果的目录
        os.makedirs(self.output_dir, exist_ok=True)
        
        # Frida的JavaScript脚本 - 专注于内存中搜索cookie
        self.jscode = """
        // 存储找到的cookie，避免重复
        var foundCookies = new Set();
        var userCookies = {};

        // 常见cookie标识符和格式
        var cookiePatterns = [
            /([\\w\\.\\-]+)=([^;]+)/g,  // 基本cookie格式
            /cookie:\\s*["']([^"']+)["']/i,  // JS中的cookie字符串
            /setCookie\\(['"]([^'"]+)['"]/i,  // JS设置cookie
            /Cookie:\\s*([^\\r\\n]+)/i  // HTTP头中的cookie
        ];

        // 用户标识符可能的模式
        var userPatterns = [
            /uid[=:]["']?([\\w\\-]+)["']?/i,
            /userId[=:]["']?([\\w\\-]+)["']?/i,
            /user_id[=:]["']?([\\w\\-]+)["']?/i,
            /username[=:]["']?([\\w\\-]+)["']?/i,
            /email[=:]["']?([^@"']+@[^@"']+)["']?/i
        ];

        // 从字符串中提取所有cookie
        function extractCookiesFromString(str) {
            var cookies = [];
            var userId = null;
            
            // 尝试识别用户
            for (var i = 0; i < userPatterns.length; i++) {
                var userMatch = str.match(userPatterns[i]);
                if (userMatch && userMatch[1]) {
                    userId = userMatch[1];
                    break;
                }
            }
            
            // 提取所有cookie
            for (var i = 0; i < cookiePatterns.length; i++) {
                var matches = str.matchAll(cookiePatterns[i]);
                for (const match of matches) {
                    var cookieStr = match[0];
                    
                    // 如果是name=value格式则单独处理
                    if (match.length > 2) {
                        cookieStr = match[1] + "=" + match[2];
                    } else if (match.length > 1) {
                        cookieStr = match[1];
                    }
                    
                    // 检查是否已找到相同cookie
                    if (!foundCookies.has(cookieStr)) {
                        foundCookies.add(cookieStr);
                        cookies.push({
                            cookie: cookieStr,
                            userId: userId || "unknown",
                            context: str.substring(Math.max(0, str.indexOf(cookieStr) - 50), 
                                                str.indexOf(cookieStr) + cookieStr.length + 50)
                        });
                    }
                }
            }
            
            return cookies;
        }

        // 搜索内存中所有可读区域
        function scanMemoryForCookies() {
            var ranges = Process.enumerateRangesSync({protection: 'r--', coalesce: true});
            send({type: "status", message: "开始扫描内存区域: " + ranges.length + "个"});
            
            var cookiesFound = 0;
            
            for (var i = 0; i < ranges.length; i++) {
                var range = ranges[i];
                if (i % 50 === 0) {
                    send({type: "progress", current: i, total: ranges.length});
                }
                
                try {
                    // 读取内存区域
                    var buffer = Memory.readByteArray(range.base, range.size);
                    
                    // 转换为字符串进行搜索
                    var str = "";
                    try {
                        // 尝试UTF-8
                        str = Memory.readUtf8String(range.base, range.size);
                    } catch (e) {
                        try {
                            // 尝试UTF-16
                            str = Memory.readUtf16String(range.base, range.size / 2);
                        } catch (e) {
                            // 跳过无法读取的内存
                            continue;
                        }
                    }
                    
                    // 寻找cookie关键词
                    if (str.toLowerCase().indexOf("cookie") !== -1 || 
                        str.indexOf("=") !== -1) {
                        
                        // 提取cookie
                        var cookies = extractCookiesFromString(str);
                        cookiesFound += cookies.length;
                        
                        if (cookies.length > 0) {
                            send({
                                type: "cookies_found",
                                address: range.base.toString(),
                                cookies: cookies
                            });
                        }
                    }
                } catch (e) {
                    // 忽略无法读取的区域
                }
            }
            
            send({type: "complete", count: cookiesFound});
        }

        // 特别处理Objective-C对象 (macOS)
        function scanObjectiveCCookies() {
            if (!ObjC.available) return;
            
            try {
                var NSHTTPCookieStorage = ObjC.classes.NSHTTPCookieStorage;
                var storage = NSHTTPCookieStorage.sharedHTTPCookieStorage();
                var cookies = storage.cookies();
                
                if (cookies) {
                    var count = cookies.count().valueOf();
                    send({type: "status", message: "发现" + count + "个NSHTTPCookie对象"});
                    
                    for (var i = 0; i < count; i++) {
                        var cookie = cookies.objectAtIndex_(i);
                        var cookieInfo = {
                            name: cookie.name().toString(),
                            value: cookie.value().toString(),
                            domain: cookie.domain().toString(),
                            path: cookie.path().toString(),
                            expires: cookie.expiresDate() ? cookie.expiresDate().toString() : "N/A",
                            secure: cookie.isSecure() ? "Yes" : "No"
                        };
                        
                        send({
                            type: "objc_cookie",
                            cookie: cookieInfo
                        });
                    }
                }
            } catch (e) {
                send({type: "error", message: "读取ObjC Cookie出错: " + e.message});
            }
        }

        // 搜索常见的cookie保存位置
        function scanCommonLocations() {
            // JavaScript变量中的cookies
            if (typeof document !== 'undefined') {
                try {
                    var docCookie = document.cookie;
                    if (docCookie && docCookie.length > 0) {
                        send({
                            type: "location_cookie",
                            source: "document.cookie",
                            value: docCookie
                        });
                        
                        var cookies = extractCookiesFromString(docCookie);
                        if (cookies.length > 0) {
                            send({
                                type: "cookies_found",
                                address: "document.cookie",
                                cookies: cookies
                            });
                        }
                    }
                } catch (e) {
                    // 忽略错误
                }
            }
            
            // 如果是Electron应用，可能会有本地存储
            try {
                var electronPaths = [
                    '/Cookies',
                    '/Cookies-journal',
                    '/Local Storage/leveldb',
                    '/Session Storage'
                ];
                
                send({type: "status", message: "检查Electron cookie存储位置"});
            } catch (e) {
                // 忽略错误
            }
        }

        // 开始扫描
        setTimeout(function() {
            send({type: "status", message: "开始分析内存中的Cookie"});
            
            // 先检查特定对象
            scanObjectiveCCookies();
            scanCommonLocations();
            
            // 然后扫描整个内存
            scanMemoryForCookies();
        }, 0);
        """

    def on_message(self, message, data):
        if message['type'] == 'send':
            payload = message['payload']
            if payload['type'] == 'status':
                print(f"[状态] {payload['message']}")
            
            elif payload['type'] == 'progress':
                percent = int((payload['current'] / payload['total']) * 100)
                print(f"[进度] {percent}% ({payload['current']}/{payload['total']})")
            
            elif payload['type'] == 'cookies_found':
                for cookie_info in payload['cookies']:
                    user_id = cookie_info['userId']
                    
                    if user_id not in self.cookies_by_user:
                        self.cookies_by_user[user_id] = set()
                    
                    self.cookies_by_user[user_id].add(cookie_info['cookie'])
                    self.total_cookies += 1
                    
                    # 打印发现的cookie
                    print(f"[Cookie] 用户: {user_id}, Cookie: {cookie_info['cookie']}")
            
            elif payload['type'] == 'objc_cookie':
                cookie_info = payload['cookie']
                cookie_str = f"{cookie_info['name']}={cookie_info['value']}"
                
                print(f"[ObjC Cookie] 名称: {cookie_info['name']}, 值: {cookie_info['value']}, 域: {cookie_info['domain']}")
                
                if 'unknown' not in self.cookies_by_user:
                    self.cookies_by_user['unknown'] = set()
                self.cookies_by_user['unknown'].add(cookie_str)
                self.total_cookies += 1
            
            elif payload['type'] == 'complete':
                print(f"\n[完成] 共找到 {payload['count']} 个cookie")
                self.save_cookies_to_file()
        
        elif message['type'] == 'error':
            print(f"[错误] {message['stack']}")

    def save_cookies_to_file(self):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{self.output_dir}/cookies_{timestamp}.json"
        
        # 转换set为list以便JSON序列化
        output = {}
        for user, cookies in self.cookies_by_user.items():
            output[user] = list(cookies)
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(output, f, indent=2, ensure_ascii=False)
        
        print(f"[保存] Cookie已保存到 {filename}")
        
        # 创建每个用户的cookie文件
        for user, cookies in self.cookies_by_user.items():
            user_filename = f"{self.output_dir}/user_{user}_{timestamp}.txt"
            with open(user_filename, 'w', encoding='utf-8') as f:
                for cookie in cookies:
                    f.write(f"{cookie}\n")
            print(f"[保存] 用户 {user} 的Cookie已保存到 {user_filename}")

    def extract_from_process(self, target_process):
        try:
            print(f"[*] 附加到进程 {target_process}")
            session = frida.attach(target_process)
            
            print("[*] 创建脚本")
            script = session.create_script(self.jscode)
            script.on('message', self.on_message)
            
            print("[*] 加载脚本")
            script.load()
            
            print("[*] 脚本正在运行，按Ctrl+C停止")
            sys.stdin.read()
        except KeyboardInterrupt:
            pass
        except Exception as e:
            print(f"[!] 错误: {e}")
        
        print("\n[*] 提取结束")
        if self.total_cookies > 0:
            print(f"[*] 共提取了 {self.total_cookies} 个cookie，分属于 {len(self.cookies_by_user)} 个用户")
            for user, cookies in self.cookies_by_user.items():
                print(f"[*] 用户 {user}: {len(cookies)} 个cookie")

def main():
    parser = argparse.ArgumentParser(description='从运行程序内存中提取Cookie')
    parser.add_argument('target', help='目标进程名称或PID')
    parser.add_argument('-o', '--output-dir', default='cookies_extracted', help='输出目录')
    parser.add_argument('-v', '--verbose', action='store_true', help='显示详细信息')
    
    args = parser.parse_args()
    
    # 解析目标进程
    target = args.target
    try:
        target = int(target)
    except ValueError:
        pass
    
    # 创建提取器
    extractor = CookieExtractor(output_dir=args.output_dir)
    
    # 执行提取
    extractor.extract_from_process(target)

if __name__ == "__main__":
    main() 